# -*- coding: utf-8 -*-
#%%
import collections
import numpy as np
from scipy import signal

#%%
BIT_RSHIFT = 16
BIT_ACT = 8
MIN_ACT, MAX_ACT = 0, np.power(2, BIT_ACT) - 1

def norm_act(img, bias, mul):
    img.shape[2] == bias.shape[0] == mul.shape[0]
    norm = (img + bias) * mul
    shif = np.right_shift(norm, BIT_RSHIFT)
    actv = np.clip(shif, MIN_ACT, MAX_ACT)
    return actv

def dwconv(img, weight, padVal=0):
    assert img.shape[2] == weight.shape[2]
    N_CH = img.shape[2]
    res = np.empty_like(img)
    for i in range(N_CH):
        res[..., i] = signal.correlate2d(
            img[..., i], weight[..., i], 
            mode='same', boundary='fill', fillvalue=padVal
        )
    return res

def pwconv(img, weight):
    shape = img.shape
    assert img.shape[2] == weight.shape[1]
    res = np.matmul(img.reshape(shape[0] * shape[1], shape[2]), weight.transpose())
    return res.reshape(shape[0], shape[1], weight.shape[0])

def dwLayer(img, weight, bias, mul, padVal=0):
    assert img.shape[2] == weight.shape[2] == bias.shape[0] == mul.shape[0]
    conv = dwconv(img, weight, padVal)
    out = norm_act(conv, bias, mul)
    return out

def pwLayer(img, weight, bias, mul):
    assert img.shape[2] == weight.shape[1]
    assert weight.shape[0] == bias.shape[0] == mul.shape[0]
    conv = pwconv(img, weight)
    out = norm_act(conv, bias, mul)
    return out

def maxPool(img):
    # kernel_size = stride = 2
    assert img.shape[0] % 2 == 0
    assert img.shape[1] % 2 == 0
    n_ch = img.shape[2]
    row, col = img.shape[0] // 2, img.shape[1] // 2
    pool = np.empty((row, col, n_ch), dtype=img.dtype)
    for i in range(n_ch):
        x = img[..., i]
        x = x.reshape(-1, col, 2).max(axis = 2)
        x = x.reshape(row, 2, -1).max(axis = 1)
        pool[..., i] = x
    return pool

def zq_resize(img):
    row_idx = np.arange(160) * 359 / 159
    row_idx = (row_idx + 0.5).astype(np.uint32)
    return img[row_idx, 1:: 2]

def find_max(img):
    assert img.shape == (20, 40, 10)
    res = np.empty(shape=(2, 7), dtype=np.int16)
    idx0 = np.unravel_index(img[..., 4].argmax(), (20, 40))
    idx1 = np.unravel_index(img[..., 9].argmax(), (20, 40))
    res[0, 0: 5] = img[idx0[0], idx0[1], 0: 5 ]
    res[0, 5: 7] = idx0[: : -1]
    res[1, 0: 5] = img[idx1[0], idx1[1], 5: 10]
    res[1, 5: 7] = idx1[: : -1]
    return res

def check(img, gold, title=""):
    assert img.shape == gold.shape
    isEqu = img == gold
    print("{}: ".format(title), end="")
    if np.all(isEqu):
        print("Pass")
    else:
        print("Boom")
        print(np.argwhere(isEqu == False))
    return np.all(isEqu)

#%%
if __name__ == "__main__":
    feature = np.load("../data/feature_int32_HWC_big.npz")
    weight = np.load("../data/skynet_weight_int32_KKC_OCIC.npz")
    val = feature["big"]
    
    # Resize
    val = zq_resize(val)
    check(val, feature["data0"], "in")
    # Bundle 1
    val = dwLayer(val, weight["L0-W"], weight["L0-B"], weight["L0-M"], 128)
    check(val, feature["conv1"], "L0")
    val = pwLayer(val, weight["L1-W"], weight["L1-B"], weight["L1-M"])
    check(val, feature["conv2"], "L1")
    val = maxPool(val)
    check(val, feature["pool1"], "P1")
    # Bundle 2
    val = dwLayer(val, weight["L2-W"], weight["L2-B"], weight["L2-M"])
    check(val, feature["conv3"], "L2")
    val = pwLayer(val, weight["L3-W"], weight["L3-B"], weight["L3-M"])
    check(val, feature["conv4"], "L3")
    val = maxPool(val)
    check(val, feature["pool2"], "P2")
    # Bundle 3
    val = dwLayer(val, weight["L4-W"], weight["L4-B"], weight["L4-M"])
    check(val, feature["conv5"], "L4")
    val = pwLayer(val, weight["L5-W"], weight["L5-B"], weight["L5-M"])
    check(val, feature["conv6"], "L5")
    
    # ByPass & ReOrder
    sp = val.shape
    # bypass = val.reshape(sp[0] // 2, 2, sp[1] // 2, 2, sp[2]).transpose((0, 2, 1, 3, 4))
    bypass = val.reshape(sp[0] // 2, 2, sp[1] // 2, 2 * sp[2]).transpose((0, 2, 1, 3))
    assert bypass.shape == (sp[0] // 2, sp[1] // 2, 2, 2 * sp[2])
    bypass = bypass.reshape(sp[0] // 2, sp[1] // 2, sp[2] * 4)
    check(bypass, feature["reorg"], "ReOrder")
    
    val = maxPool(val)
    check(val, feature["pool3"], "P3")
    # Bundle 4
    val = dwLayer(val, weight["L6-W"], weight["L6-B"], weight["L6-M"])
    check(val, feature["conv7"], "L6")
    val = pwLayer(val, weight["L7-W"], weight["L7-B"], weight["L7-M"])
    check(val, feature["conv8"], "L7")
    # Bundle 5
    val = dwLayer(val, weight["L8-W"], weight["L8-B"], weight["L8-M"])
    check(val, feature["conv9"], "L8")
    val = pwLayer(val, weight["L9-W"], weight["L9-B"], weight["L9-M"])
    check(val, feature["conv10"], "L9")
    
    # concatenate
    val = np.concatenate((bypass, val), axis=2)
    check(val, feature["cat"], "Concat")
    
    # Bundle 6
    val = dwLayer(val, weight["L10-W"], weight["L10-B"], weight["L10-M"])
    check(val, feature["conv11"], "L10")
    val = pwLayer(val, weight["L11-W"], weight["L11-B"], weight["L11-M"])
    check(val, feature["conv12"], "L11")
    # Output
    val = pwconv(val, weight["L12-W"])
    check(val, feature["conv13"], "L12")
    # findmax
    val = find_max(val)
    val = val.reshape(1, *val.shape)
    check(val, feature["fmax"], "out")

